package com.borened.common.boot.interceptor;

import cn.hutool.extra.servlet.ServletUtil;
import cn.hutool.json.JSONUtil;
import com.borened.common.core.base.result.R;
import com.borened.common.core.base.result.ResultCode;
import com.borened.common.core.security.RequestRateLimit;
import com.borened.common.core.util.ip.IpUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

@Slf4j
public class RequestRateLimitInterceptor implements HandlerInterceptor {

    public static final String REQ_LIMIT_KEY = "req_rate_limit_";

    @Autowired(required = false)
    private RedisTemplate<String, String> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }
        HandlerMethod handlerMethod = (HandlerMethod) handler;
        //方法注解
        RequestRateLimit methodAnnotation = handlerMethod.getMethodAnnotation(RequestRateLimit.class);
        //类注解
        RequestRateLimit classAnnotation = handlerMethod.getBean().getClass().getAnnotation(RequestRateLimit.class);
        boolean checkResult = true;
        if (methodAnnotation != null) {
            checkResult = validateRequestLimit(request, methodAnnotation.count(), methodAnnotation.time(), methodAnnotation.timeUnit());
        } else if (classAnnotation != null) {
            checkResult = validateRequestLimit(request, classAnnotation.count(), classAnnotation.time(), methodAnnotation.timeUnit());
        }
        if (checkResult) {
            return true;
        } else {
            ServletUtil.write(response, JSONUtil.toJsonStr(R.fail(ResultCode.REQ_RATE_LIMIT)), "application/json;charset=utf-8");
            return false;
        }
    }

    /**
     * 验证接口的访问频次限制
     *
     * @param request
     * @return
     */
    private boolean validateRequestLimit(HttpServletRequest request, int maxSize, long timeOut, TimeUnit timeUnit) {
        boolean pass = true;
        try {
            String ip = IpUtils.getIpAddr(request);
            String url = request.getRequestURL().toString();
            String key = REQ_LIMIT_KEY.concat(url).concat(ip);
            long count = redisTemplate.opsForValue().increment(key, 1);
            if (count == 1) {
                redisTemplate.expire(key, timeOut, timeUnit);
            }
            if (count > maxSize) {
                log.info("用户IP[" + ip + "]访问地址[" + url + "]超过了限定的次数[" + maxSize + "]");
                pass = false;
            }
        } catch (Exception e) {
            log.error("发生异常: ", e);
        }
        return pass;
    }

}

